/*
 * Copyright (c) 2025.
 * All rights reserved.
 */

#include "driver.h"
#include "ioctl_communicator.h"
#include "string_utils.h"

#define NAME_BUFFER_SIZE        128
#define PREFIX_SEND_DEVICE      L"\\Device\\"IOCTL_SEND_DEV_NAME
#define PREFIX_SEND_SYM_LINK    L"\\DosDevices\\"IOCTL_SEND_DEV_NAME
#define PREFIX_RECV_DEVICE      L"\\Device\\"IOCTL_RECV_DEV_NAME
#define PREFIX_RECV_SYM_LINK    L"\\DosDevices\\"IOCTL_RECV_DEV_NAME

#define MSG_ITEM_SIZE           5
#define STRING_NAME_TAG         'name'
#define NAME_INDEX_BUF_SIZE     16
#define NAME_INDEX_BASE         10

#define DEVICE_ID_SEND          1
#define DEVICE_ID_RECV          2

#define MESSAGE_NO_REPLY        (-1)

typedef struct _DEVICE_EXTENSION {
    UCHAR DeviceID;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

UNICODE_STRING  g_IoSendSymLinkName;
UNICODE_STRING  g_IoSendDeviceName;
UNICODE_STRING  g_IoRecvSymLinkName;
UNICODE_STRING  g_IoRecvDeviceName;
PDEVICE_OBJECT  g_SendDevice = NULL;
PDEVICE_OBJECT  g_RecvDevice = NULL;

PMSG_RECV_FUNC  g_MsgAsyncFunc;
PIRP            g_IoctlPendingIrp = NULL;
KSPIN_LOCK      g_MsgLock;
MSG_ITEM        g_Messages[MSG_ITEM_SIZE];

VOID IoctlUnload(
    _In_ PDRIVER_OBJECT driverObj)
{
    UNREFERENCED_PARAMETER(driverObj);

    if (g_IoctlPendingIrp) {
        IoSetCancelRoutine(g_IoctlPendingIrp, NULL);
        IoCompleteRequest(g_IoctlPendingIrp, IO_NO_INCREMENT);  // clear pended irp
        g_IoctlPendingIrp = NULL;
    }

    if (g_SendDevice != NULL) {
        IoDeleteDevice(g_SendDevice);
        DBG_PRINT_IO("Delete send device,%wZ\n", g_IoSendDeviceName);
    }
    if (g_RecvDevice != NULL) {
        IoDeleteDevice(g_RecvDevice);
        DBG_PRINT_IO("Delete receive device,%wZ\n", g_IoRecvDeviceName);
    }
    if (STATUS_SUCCESS == IoDeleteSymbolicLink(&g_IoSendSymLinkName)) {
        DBG_PRINT_IO("Delete send symbolic link,%wZ\n", g_IoSendSymLinkName);
    }
    if (STATUS_SUCCESS == IoDeleteSymbolicLink(&g_IoRecvSymLinkName)) {
        DBG_PRINT_IO("Delete receive symbolic link,%wZ\n", g_IoRecvSymLinkName);
    }

    RtlFreeUnicodeString(&g_IoSendDeviceName);
    RtlFreeUnicodeString(&g_IoSendSymLinkName);
    RtlFreeUnicodeString(&g_IoRecvDeviceName);
    RtlFreeUnicodeString(&g_IoRecvSymLinkName);
}
VOID SetMsgState(
    _In_ PMSG_ITEM pMsg,
    _In_ CHAR state)
{
    if (pMsg == NULL) {
        return;
    }
    pMsg->State = state;
}

/// <summary>
/// Init callback and local state
/// </summary>
/// <param name="driverObj"></param>
/// <param name="msgRecvFunc"></param>
/// <returns></returns>
NTSTATUS IoctlLoad(
    _In_ PDRIVER_OBJECT driverObj,
    _In_ PMSG_RECV_FUNC msgRecvFunc)
{
    if (driverObj == NULL) {
        DBG_PRINT_IO("Init Fail,driver object is null\n");
        return STATUS_INVALID_PARAMETER;
    }
    if (msgRecvFunc == NULL) {
        DBG_PRINT_IO("Init Fail,receive function is null\n");
        return STATUS_INVALID_PARAMETER;
    }

    // routine for handling IO requests from userland
    driverObj->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IoctlHandleMsg;

    // routines that will execute once a handle to our device's symbolik link is opened/closed
    driverObj->MajorFunction[IRP_MJ_CREATE]  = IoctlConnect;
    driverObj->MajorFunction[IRP_MJ_CLOSE]   = IoctlDisconnect;

    g_MsgAsyncFunc = msgRecvFunc;
    KeInitializeSpinLock(&g_MsgLock);
    // init string
    if (!InitUnicodeString(&g_IoSendDeviceName, NAME_BUFFER_SIZE, STRING_NAME_TAG) ||
        !InitUnicodeString(&g_IoSendSymLinkName, NAME_BUFFER_SIZE, STRING_NAME_TAG) ||
        !InitUnicodeString(&g_IoRecvSymLinkName, NAME_BUFFER_SIZE, STRING_NAME_TAG) ||
        !InitUnicodeString(&g_IoRecvDeviceName, NAME_BUFFER_SIZE, STRING_NAME_TAG)) {
        DBG_PRINT("InitNameString: Fail\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    // init message list
    for (CHAR i = 0; i < MSG_ITEM_SIZE; i++) {
        g_Messages[i].Id = i;
        g_Messages[i].Message = 0;
        SetMsgState(&g_Messages[i], MSG_STATE_IDLE);
        KeInitializeEvent(&g_Messages[i].Event, NotificationEvent, FALSE);
    }

    NTSTATUS status;
    status = IoctlCreateDevice(driverObj);
    if (STATUS_SUCCESS != status) {
        DBG_PRINT("IoctlCreateDevice: Fail\n");
        return status;
    }
    return status;
}
PMSG_ITEM GetIdleMessage(VOID)
{
    for (int i = 0; i < MSG_ITEM_SIZE; i++) {
        if (g_Messages[i].State == MSG_STATE_IDLE) {
			return &g_Messages[i];
        }
    }
    DBG_PRINT("GetIdleMessage: Fail\n");
    return NULL;
}

NTSTATUS BuildIndexName(
    int index,
    PUNICODE_STRING pNameStr,
    PCWSTR namePrefix)
{
    NTSTATUS        status;
    UNICODE_STRING  indexStr;
    WCHAR           indexBuf[NAME_INDEX_BUF_SIZE] = { 0 };

    RtlInitEmptyUnicodeString(&indexStr, indexBuf, NAME_INDEX_BUF_SIZE * sizeof(WCHAR));
    RtlIntegerToUnicodeString(index, NAME_INDEX_BASE, &indexStr);
    RtlInitEmptyUnicodeString(pNameStr, pNameStr->Buffer, pNameStr->MaximumLength);
    status = RtlAppendUnicodeToString(pNameStr, namePrefix);
    if (!NT_SUCCESS(status)) {
        DBG_PRINT_IO("Failed to append name,%d,%d\n", index, status);
        return status;
    }
    status = RtlAppendUnicodeStringToString(pNameStr, &indexStr);
    if (!NT_SUCCESS(status)) {
        DBG_PRINT_IO("Failed to append index,%d,%d\n", index, status);
        return status;
    }
    return STATUS_SUCCESS;
}

NTSTATUS CreateDevice(
    _In_ PDRIVER_OBJECT driverObj,
    _In_ PUNICODE_STRING pNameStr,
    _In_ PCWSTR namePrefix,
    _Out_ PDEVICE_OBJECT *pDevOut)
{
    if (driverObj == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    NTSTATUS status;

    for (int i = 0; i < IOCTL_MAX_DEVICE_SIZE; i++) {
        status = BuildIndexName(i, pNameStr, namePrefix);
        if (!NT_SUCCESS(status)) {
            DBG_PRINT_IO("BuildIndexName fail,%d\n", i);
            return status;
        }
        status = IoCreateDevice(
            driverObj,
            sizeof(DEVICE_EXTENSION), // Device extension is not needed.
            pNameStr,
            FILE_DEVICE_UNKNOWN,
            0, //FILE_DEVICE_SECURE_OPEN,
            FALSE, // Multiple handles can access the device at the same time
            pDevOut);
        if (NT_SUCCESS(status)) {
            break;
        }
        if (status == STATUS_OBJECT_NAME_COLLISION) {
            DBG_PRINT_IO("Device Name conflict,%wZ\n", *pNameStr);
            continue;
        }
    }

    if (!NT_SUCCESS(status)) {
        DBG_PRINT_IO("Create Device,%wZ, Fail\n", *pNameStr);
        return status;
    }
    if (*pDevOut != NULL) {
        (*pDevOut)->Flags |= DO_BUFFERED_IO; //buffer mode
    }
    return status;
}
NTSTATUS CreateSymbolicLink(
    _In_ PDRIVER_OBJECT driverObj,
    _In_ PUNICODE_STRING pLinkName,
    _In_ PUNICODE_STRING pDeviceName,
    _In_ PCWSTR namePrefix)
{
    if (driverObj == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    NTSTATUS status;

    for (int i = 0; i < IOCTL_MAX_DEVICE_SIZE; i++) {
        status = BuildIndexName(i, pLinkName, namePrefix);
        if (!NT_SUCCESS(status)) {
            DBG_PRINT_IO("BuildIndexName fail,%d\n", status);
            return status;
        }
        status = IoCreateSymbolicLink(pLinkName, pDeviceName);
        if (NT_SUCCESS(status)) {
            break;
        }
        if (status == STATUS_OBJECT_NAME_COLLISION) {
            DBG_PRINT_IO("SymbolicLink Name conflict,%wZ\n", *pLinkName);
            continue;
        }
    }

    if (!NT_SUCCESS(status)) {
        DBG_PRINT_IO("Create Symbolic,%wZ, Fail\n", *pLinkName);
        return status;
    }
    return status;
}
NTSTATUS IoctlCreateDevice(
    _In_ PDRIVER_OBJECT driverObj)
{
    if (driverObj == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    NTSTATUS status;
    PDEVICE_EXTENSION pDevExt;
   
    // create send device
    status = CreateDevice(driverObj, &g_IoSendDeviceName, PREFIX_SEND_DEVICE, &g_SendDevice);
    if (!NT_SUCCESS(status) || g_SendDevice == NULL) {
        return status;
    }
    DBG_PRINT_IO("Create Device,%wZ, Succeed\n", g_IoSendDeviceName);
    pDevExt = (PDEVICE_EXTENSION)(g_SendDevice->DeviceExtension);
    if (pDevExt != NULL) {
        pDevExt->DeviceID = DEVICE_ID_SEND;
    }

    status = CreateSymbolicLink(driverObj, &g_IoSendSymLinkName, &g_IoSendDeviceName, PREFIX_SEND_SYM_LINK);
    if (!NT_SUCCESS(status)) {
        IoDeleteDevice(g_SendDevice);
        return status;
    }
    DBG_PRINT_IO("Create Symbolic,%wZ, Succeed\n", g_IoSendSymLinkName);
    
    // create receive device
    status = CreateDevice(driverObj, &g_IoRecvDeviceName, PREFIX_RECV_DEVICE, &g_RecvDevice);
    if (!NT_SUCCESS(status)) {
        return status;
    }
    DBG_PRINT_IO("Create Device,%wZ, Succeed\n", g_IoRecvDeviceName);
    pDevExt = (PDEVICE_EXTENSION)(g_RecvDevice->DeviceExtension);
    if (pDevExt != NULL) {
        pDevExt->DeviceID = DEVICE_ID_RECV;
    }

    status = CreateSymbolicLink(driverObj, &g_IoRecvSymLinkName, &g_IoRecvDeviceName, PREFIX_RECV_SYM_LINK);
    if (!NT_SUCCESS(status)) {
        IoDeleteDevice(g_RecvDevice);
        return status;
    }
    DBG_PRINT_IO("Create Symbolic,%wZ, Succeed\n", g_IoRecvSymLinkName);

    return status;
}


/// <summary>
/// Complete pended message request.
/// </summary>
/// <returns></returns>
BOOLEAN CompletePendingIrp(
    _In_ ULONG msgSize)
{
    if (g_IoctlPendingIrp == NULL) {
        return FALSE;
    }
    
    g_IoctlPendingIrp->IoStatus.Information = msgSize;
    g_IoctlPendingIrp->IoStatus.Status = STATUS_SUCCESS;
    IoSetCancelRoutine(g_IoctlPendingIrp, NULL);
    IoCompleteRequest(g_IoctlPendingIrp, IO_NO_INCREMENT);
    g_IoctlPendingIrp = NULL;
    return TRUE;
}

VOID ReleaseAllMessageWaiting()
{
    for (CHAR i = 0; i < MSG_ITEM_SIZE; i++) {
        KeSetEvent(&g_Messages[i].Event, 0, FALSE);
        SetMsgState(&g_Messages[i], MSG_STATE_IDLE);
    }
}

NTSTATUS IoctlConnect(
    _In_ PDEVICE_OBJECT deviceObj,
    _In_ PIRP irp)
{
    UNREFERENCED_PARAMETER(deviceObj);

    if (irp == NULL) {
        return STATUS_INVALID_PARAMETER;
    }
    PDEVICE_EXTENSION pDevExt;
    pDevExt = (PDEVICE_EXTENSION)(deviceObj->DeviceExtension);
    if (pDevExt != NULL) {
        if (pDevExt->DeviceID == DEVICE_ID_SEND) {
            DBG_PRINT_IO("Connect,%wZ\n", g_IoSendSymLinkName);
        } else if (pDevExt->DeviceID == DEVICE_ID_RECV) {
            DBG_PRINT_IO("Connect,%wZ\n", g_IoRecvSymLinkName);
        }
    }

    // release current pending irp
    if (CompletePendingIrp(0)) {
        DBG_PRINT_IO("Connect,clear pending irp\n");
    }

    irp->IoStatus.Information = 0;
    irp->IoStatus.Status = STATUS_SUCCESS;
    IoSetCancelRoutine(irp, NULL);
    IoCompleteRequest(irp, IO_NO_INCREMENT);

    return STATUS_SUCCESS;
}

NTSTATUS IoctlDisconnect(
    _In_ PDEVICE_OBJECT deviceObj,
    _In_ PIRP irp)
{
    UNREFERENCED_PARAMETER(deviceObj);

    if (irp == NULL) {
        return STATUS_INVALID_PARAMETER;
    }
    PDEVICE_EXTENSION pDevExt;
    pDevExt = (PDEVICE_EXTENSION)(deviceObj->DeviceExtension);
    if (pDevExt != NULL) {
        if (pDevExt->DeviceID == DEVICE_ID_SEND) {
            DBG_PRINT_IO("Disconnect,%wZ\n", g_IoSendSymLinkName);
        } else if (pDevExt->DeviceID == DEVICE_ID_RECV) {
            DBG_PRINT_IO("Disconnect,%wZ\n", g_IoRecvSymLinkName);
        }
    }

    // release current pending irp
    if (CompletePendingIrp(0)) {
        DBG_PRINT_IO("Disconnect,clear pending irp\n");
    }
    // release message waiting
    ReleaseAllMessageWaiting();

    irp->IoStatus.Information = 0;
    irp->IoStatus.Status = STATUS_SUCCESS;
    IoSetCancelRoutine(irp, NULL);
    IoCompleteRequest(irp, IO_NO_INCREMENT);
        
    return STATUS_SUCCESS;
}

VOID IoctlCancelIrp(
    _In_ PDEVICE_OBJECT deviceObj,
    _In_ PIRP irp)
{
    UNREFERENCED_PARAMETER(deviceObj);
    DBG_PRINT_IO("CancelIrp\n");

    // release message waiting
    ReleaseAllMessageWaiting();

    if (irp == g_IoctlPendingIrp)
        g_IoctlPendingIrp = NULL;

    if (irp != NULL) {
        irp->IoStatus.Status = STATUS_CANCELLED;
        irp->IoStatus.Information = 0;
        IoSetCancelRoutine(irp, NULL);
        IoCompleteRequest(irp, IO_NO_INCREMENT);

        // Important
        IoReleaseCancelSpinLock(irp->CancelIrql);
    }
}
/// <summary>
/// Triggered by the user layer, no reentry exists.
/// </summary>
/// <param name="deviceObject"></param>
/// <param name="Irp"></param>
/// <returns></returns>
NTSTATUS IoctlHandleMsg(
    _In_ PDEVICE_OBJECT deviceObj,
    _In_ PIRP irp)
{
    UNREFERENCED_PARAMETER(deviceObj);
    if (irp == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    PIO_STACK_LOCATION  stack;
    NTSTATUS            status = STATUS_SUCCESS;
    ULONG               msgCode;
    PIO_MATCH_MSG       msgBuf = NULL;
    ULONG               msgCapacity = 0;

    stack = IoGetCurrentIrpStackLocation(irp);
    msgBuf = (PIO_MATCH_MSG)irp->AssociatedIrp.SystemBuffer;
    if (stack == NULL || msgBuf == NULL) {
        return STATUS_INVALID_PARAMETER;
    }
    msgCode = stack->Parameters.DeviceIoControl.IoControlCode;
    msgCapacity = stack->Parameters.DeviceIoControl.OutputBufferLength;

    if (msgCode == IOCTL_MSG_ASYNC) {
        g_MsgAsyncFunc(msgBuf->Message);
        if (g_IoctlPendingIrp == NULL) {
            irp->IoStatus.Status = STATUS_PENDING;
            irp->IoStatus.Information = 0;
            IoMarkIrpPending(irp);  // set next irp
            IoSetCancelRoutine(irp, IoctlCancelIrp);
            g_IoctlPendingIrp = irp;
            return STATUS_PENDING;
        } else { //error
            status = STATUS_INVALID_PARAMETER;
        }
    } else if (msgCode == IOCTL_MSG_SYNC) {
        if (msgBuf->Id >= 0 && msgBuf->Id < MSG_ITEM_SIZE) {
            g_Messages[msgBuf->Id].Message = msgBuf->Message;
            KeSetEvent(&g_Messages[msgBuf->Id].Event, 0, FALSE);
        }
        msgBuf->Message = MSG_STATE_DEVICE_OK;
        status = STATUS_SUCCESS;
    } else {
        status = STATUS_INVALID_DEVICE_REQUEST;
    }
    irp->IoStatus.Status = status;
    irp->IoStatus.Information = (status == STATUS_SUCCESS) ? msgCapacity : 0;
    IoSetCancelRoutine(irp, NULL);
    IoCompleteRequest(irp, IO_NO_INCREMENT);
    
    return status;
}

/// <summary>
/// return the size of message buffer
/// </summary>
/// <param name="msgBuf"></param>
/// <param name="msgFlag"></param>
/// <param name="msgParam"></param>
/// <returns></returns>
ULONG BuildMsgBuffer(
    _In_ PIO_MATCH_MSG msgBuf,
    _In_ MSG_PARAM msgParam)
{
    if (msgBuf == NULL || msgBuf->NameMaxLength == 0 || g_IoctlPendingIrp == NULL) {
        return 0;
    }
    PIO_STACK_LOCATION  stack = IoGetCurrentIrpStackLocation(g_IoctlPendingIrp);
    if (stack == NULL) {
        return 0;
    }
    USHORT outLen = min(msgBuf->NameMaxLength, msgParam.filePath->Length);
    if (outLen == 0) {
        return 0;
    }
    if (memcpy_s(msgBuf->FileName, outLen, msgParam.filePath->Buffer, outLen) != 0) {
        return 0;
    }
    msgBuf->CurrentPid = msgParam.pid;
    msgBuf->ParentPid = msgParam.parentPid;
    msgBuf->Message = msgParam.type;
    msgBuf->NameMaxLength = outLen;
    return stack->Parameters.DeviceIoControl.OutputBufferLength;
}

/// <summary>
/// Will Send and wait reply.
/// High frequency concurrency function
/// </summary>
/// <returns></returns>
BOOLEAN IoctlSendMsg(
    _In_ MSG_PARAM msgParam,
    _In_ BOOLEAN waitReply,
    _Out_ PMSG_ITEM* msgItem)
{
    *msgItem = NULL;
    if (g_IoctlPendingIrp == NULL || IsNullOrEmpty(msgParam.filePath)) {
        return FALSE;
    }
    KIRQL OldIrql;

    KeAcquireSpinLock(&g_MsgLock, &OldIrql); //======================Lock PendingIrp Start=============
    if (g_IoctlPendingIrp == NULL) {
        KeReleaseSpinLock(&g_MsgLock, OldIrql);
        return FALSE;
    }
    PMSG_ITEM       idleMsg = GetIdleMessage();
    PIO_MATCH_MSG   msgBuf = (PIO_MATCH_MSG)g_IoctlPendingIrp->AssociatedIrp.SystemBuffer;
    ULONG           msgCapacity = BuildMsgBuffer(msgBuf, msgParam);

    if (idleMsg == NULL || msgBuf == NULL || msgCapacity == 0) {
        KeReleaseSpinLock(&g_MsgLock, OldIrql);
        return FALSE;
    }

    if (waitReply) {
        msgBuf->Id = idleMsg->Id;
        SetMsgState(idleMsg, MSG_STATE_BUSY);
    } else {
        msgBuf->Id = MESSAGE_NO_REPLY;
    }
    CompletePendingIrp(msgCapacity); // set Ioctl_PendingIrp = null.
    KeReleaseSpinLock(&g_MsgLock, OldIrql); //======================Lock PendingIrp End=============

    if (!waitReply) {
        return TRUE;
    }
    KeResetEvent(&idleMsg->Event);
    if (STATUS_SUCCESS == KeWaitForSingleObject(&idleMsg->Event, Executive, KernelMode, FALSE, NULL)) {
        *msgItem = idleMsg;
        return TRUE;
    }
    return FALSE;
}